<template>
  <div class="custom-tree">
    <div v-for="node in flatTreeData" :key="node.id" class="tree-node" :style="{ marginLeft: node.level * 20 + 'px' }">
      <div class="tree-node-label" @click="handleNodeClick(node)">
        <el-icon v-if="node.children && node.children.length > 0">
          <CaretRight :class="!node.expanded ? 'icon-caret-right' : 'icon-caret-bottom'" />
        </el-icon>
        <el-checkbox v-if="avaChecked" v-model="node.checked" @change="handleCheckChange(node)"
          :disabled="node.disabled" :indeterminate="node.indeterminate"></el-checkbox>
        <span class="tree-node-label-text">{{ node.label }}</span>
      </div>

      <!-- 添加自定义内容 -->
      <div v-if="node.isLast" class="tree-node-custom-content">
        <slot name="treeOperation" :node="node">
          <!-- 在父组件中使用，比如el-form -->
        </slot>
      </div>
    </div>
  </div>
</template>

<script lang="ts" setup>
import { ref, reactive, toRefs, watch, PropType } from 'vue';
import { ElCheckbox } from 'element-plus';
import { CaretRight } from '@element-plus/icons-vue';

interface TreeNode {
  id: number;
  label: string;
  expanded: boolean;
  checked: boolean;
  disabled: boolean;
  indeterminate: boolean;
  level: number; // 新增属性，表示节点层级
  children?: TreeNode[];
  parentId: number;
  isLast: boolean;
}

const props = defineProps({
  treeData: {
    type: Array as PropType<TreeNode[]>,
    required: true,
  },
  avaChecked: {
    type: Boolean,
    default: false,
  },
  alleExpanded: {
    type: Boolean,
    default: false,
  },
  allChecked: {
    type: Boolean,
    default: false,
  },
});

const flatTreeData = ref<TreeNode[]>([]);

const flattenTree = (nodes: TreeNode[], level: number = 0) => {
  nodes.forEach(node => {
    node.level = level;
    flatTreeData.value.push(node);
    if (node.expanded && node.children && node.children.length > 0) {
      flattenTree(node.children, level + 1);
    }
  });
};

const initializeTree = () => {
  flatTreeData.value = [];
  flattenTree(props.treeData);
};

watch(() => props.treeData, initializeTree, { immediate: true });

const emit = defineEmits(['node-checked', 'node-click']);

const handleNodeClick = (node: TreeNode) => {
  node.expanded = !node.expanded;
  initializeTree();
  emit('node-click', node);
};

const handleCheckChange = (node: TreeNode) => {
  updateChildCheckStatus(node);
  updateParentCheckStatus(node);
  emit('node-checked', node);
};

const updateChildCheckStatus = (node: TreeNode) => {
  if (node.children) {
    node.children.forEach(child => {
      child.checked = node.checked;
      child.indeterminate = false; // 清除子节点的 indeterminate 状态
      updateChildCheckStatus(child);
    });
  }
};

const updateParentCheckStatus = (node: TreeNode) => {
  const parentNode = findParentNode(node);
  if (parentNode) {
    const allChecked = parentNode.children?.every(child => child.checked) || false;
    const someChecked = parentNode.children?.some(child => child.checked || child.indeterminate) || false;

    parentNode.checked = allChecked;
    parentNode.indeterminate = someChecked && !allChecked;

    updateParentCheckStatus(parentNode);
  }
};

const findParentNode = (node: TreeNode): TreeNode | undefined => {
  const findNodeById = (nodes: TreeNode[], id: number): TreeNode | undefined => {
    for (const n of nodes) {
      if (n.id === id) return n;
      if (n.children) {
        const found = findNodeById(n.children, id);
        if (found) return found;
      }
    }
    return undefined;
  };

  return findNodeById(props.treeData, node.parentId);
};

const selectAll = () => {
  if (props.allChecked) {
    selectNode(props.treeData);
  }
};

const selectNode = (nodes: TreeNode[]) => {
  nodes.forEach(node => {
    node.checked = true;
    if (node.children) {
      selectNode(node.children);
    }
  });
};

const unSelectAll = () => {
  if (!props.allChecked) {
    unSelectNode(props.treeData);
  }
};

const unSelectNode = (nodes: TreeNode[]) => {
  nodes.forEach(node => {
    node.checked = false;
    if (node.children) {
      unSelectNode(node.children);
    }
  });
};

watch(() => props.allChecked, () => {
  selectAll();
  unSelectAll();
});

const expandAll = () => {
  if (props.alleExpanded) {
    expandNode(props.treeData);
  }
};

const expandNode = (nodes: TreeNode[]) => {
  nodes.forEach(node => {
    node.expanded = true;
    if (node.children) {
      expandNode(node.children);
    }
  });
};

const collapseAll = () => {
  if (!props.alleExpanded) {
    collapseNode(props.treeData);
  }
};

const collapseNode = (nodes: TreeNode[]) => {
  nodes.forEach(node => {
    node.expanded = false;
    if (node.children) {
      collapseNode(node.children);
    }
  });
};

watch(() => props.alleExpanded, () => {
  expandAll();
  collapseAll();
});

defineExpose({
  selectAll,
  expandAll,
  unSelectAll,
  collapseAll,
});
</script>

<style scoped>
.custom-tree {
  width: 100%;
}

.tree-node {
  margin-left: 1rem;
}

.tree-node-label {
  display: flex;
  align-items: center;
  cursor: pointer;
}

.tree-node-custom-content {
  margin-left: 20px;
}

.tree-node-label-text {
  margin-left: 10px;
}

.icon-caret-right {
  color: #acafb6;
  margin-right: 4px;
}

.icon-caret-bottom {
  color: #acafb6;
  transform: rotate(90deg);
  margin-right: 4px;
}
</style>
